import numpy as np
import scipy.io.wavfile as siowav
from numpy.lib import stride_tricks
import matplotlib.pyplot as plt


# ============= Window function =================================

def rect_window(n):
    return np.ones(n)


def gaussian_window(n):
    sigma = 0.5
    power = -0.5 * ((np.arange(n) - (n - 1)/2)/(sigma * (n - 1)/2))**2
    return np.exp(power)


def hamming_window(n):
    return 0.54 - 0.46 * np.cos(2 * np.pi * np.arange(n)/(n - 1))


def hann_window(n):
    return 0.5 * (1 - np.cos(2 * np.pi * np.arange(n)/(n - 1)))


# ===== Calculate short time energy and cross zero rate ==============

def get_frames(wav_file, window_len, hop_len, is_norm=False):
    sr, signals = siowav.read(wav_file)
    frame_size = int(window_len * sr)
    hop_size = int(hop_len * sr)
    frame_num = np.floor((len(signals) - frame_size) / float(hop_size)).astype(int) + 1
    ele_stride = signals.strides[-1]
    frames = stride_tricks.as_strided(signals, shape=(frame_num, frame_size),
                                      strides=(ele_stride * hop_size, ele_stride)).copy()
    frames = frames.astype(float)
    if is_norm:
        frames = frames/np.max(frames)
    return frames


def st_energy(wav_file, window_len, hop_len, window_func, is_log=True, is_norm=True):

    frames = get_frames(wav_file, window_len, hop_len)
    frame_num, frame_size = frames.shape
    wined_data = np.multiply(frames, window_func(frame_size))
    energy = np.sum(wined_data**2, axis=1)
    if is_log:
        energy = np.log(energy + 1e-6)
    if is_norm:
        energy = (energy - np.min(energy))/(np.max(energy) - np.min(energy))
    t = window_len / 2 + np.arange(frame_num) * hop_len
    return energy, t


def st_zcr(wav_file, window_len, hop_len, threshold, window_func, is_norm=True):
    frames = get_frames(wav_file, window_len, hop_len)
    frame_num, frame_size = frames.shape
    assign1 = np.where((frames[:, :-1] - threshold) * (frames[:, 1:] - threshold) < 0, 1, 0)
    assign2 = np.where((frames[:, :-1] + threshold) * (frames[:, 1:] + threshold) < 0, 1, 0)
    wined_data = window_func(frame_size - 1)
    zcr = np.sum(wined_data * (assign1 + assign2), axis=1)
    if is_norm:
        zcr = (zcr - np.min(zcr))/(np.max(zcr) - np.min(zcr))
    t = window_len / 2 + np.arange(frame_num) * hop_len
    return zcr, t


# ============== 双门限法 ===================


def get_above_mh(energy_arr, mh):
    status = 0
    point = np.array([0, 0])
    points = list()
    for i, energy in zip(range(len(energy_arr)), energy_arr):
        if energy > mh and status == 0:
            point[0] = i
            status = 1
        elif energy <= mh and status == 1:
            point[1] = i - 1
            status = 0
            points.append(point.copy())
    return points


def expand_above_ml(energy_arr, ml, points):
    for point in points:
        i = point[0]
        j = point[1]
        while i >= 0 and energy_arr[i] >= ml:
            point[0] = i
            i -= 1
        while j < len(energy_arr) and energy_arr[j] >= ml:
            point[1] = j
            j += 1
    return points


def expand_above_zm(czr_arr, zm, points, limit_num):
    for point in points:
        l_num = limit_num
        i = point[0]
        while i >= 0 and czr_arr[i] >= zm and l_num >= 0:
            point[0] = i
            i -= 1
            l_num -= 1
    return points


def trim_points(points):
    point_n = np.array([0, 0])
    points_n = list()
    last_max = 0
    for point in points:
        if last_max > 0 and last_max >= point[0]:
            p_last = points_n.pop()
            point_n[0] = min(p_last[0], point[0])
            point_n[1] = max(p_last[1], point[1])
            points_n.append(point_n.copy())
            last_max = point_n[1]
        else:
            last_max = point[1]
            points_n.append(point)
    return points_n


def get_above_ze(zcr_arr, ze):
    status = 0
    point = np.array([0, 0])
    points = list()
    for i, zcr in zip(range(len(zcr_arr)), zcr_arr):
        if zcr > ze and status == 0:
            point[0] = i
            status = 1
        elif zcr <= ze and status == 1:
            point[1] = i - 1
            status = 0
            points.append(point.copy())
    return points


def split_points(e_points, z_points, min_part_len=3):
    points = list()
    point = np.array([0, 0])
    for e_point in e_points:
        for z_point in z_points:
            if z_point[0] > e_point[0] + min_part_len and z_point[1] < e_point[1] - min_part_len:
                point[0] = e_point[0]
                point[1] = z_point[0] - 2
                e_point[0] = z_point[0]
                points.append(point.copy())
        points.append(e_point)
    return points


def draw(sr, signals, t, energy, czr, points):
    plt.subplot(3, 1, 1)
    plt.plot(1 / sr * np.arange(len(signals)), signals, color='r')
    plt.ylabel('wave')
    for point in points:
        plt.axvline(t[point[0]], color='b')
        plt.axvline(t[point[1]], color='g')

    plt.subplot(3, 1, 2)
    plt.plot(t, energy, color='r')
    plt.ylabel('energy')

    plt.subplot(3, 1, 3)
    plt.plot(t, czr, color='r')
    plt.ylabel('zero-crossing rate')
    plt.xlabel('t')
    plt.show()


def draw_a_canvas(sr, signals, t, energy, czr, points, ml, zm):
    plt.axhline(ml, color='k')
    plt.axhline(zm, color='k')
    plt.plot(1 / sr * np.arange(len(signals)), signals, color='b')
    for point in points:
        print(t[point[0]], t[point[1]])
        # print(t[point[1]])
        plt.axvline(t[point[0]], color='g')
        plt.axvline(t[point[1]], color='r')
    plt.plot(t, energy, color='c')
    plt.plot(t, czr, color='m')
    plt.show()


def draw_only_wav(sr, signals, t, energy, czr, points, title=None):
    plt.plot(1 / sr * np.arange(len(signals)), signals, color='r')
    for point in points:
        print(t[point[0]], t[point[1]])
        # print(t[point[1]])
        plt.axvline(t[point[0]], color='b')
        plt.axvline(t[point[1]], color='g')
    if not (title is None):
        plt.title(title)
    plt.xlabel('t')
    # plt.savefig()
    plt.show()


# 双门限法
def vad():
    mh = 0.6
    ml = 0.4
    zm = 0.02
    zm_expand_limit = 0.08
    wav_file = '/Users/d/Project/python_project/speech/wavs/isolated.wav'
    window_len = 0.03
    hop_len = 0.005
    zcr_threshold = 100
    sr, signals = siowav.read(wav_file)
    signals = signals/np.max(signals)
    energy, t = st_energy(wav_file, window_len, hop_len, hamming_window)
    czr, t = st_zcr(wav_file, window_len, hop_len, zcr_threshold, hamming_window)
    points = get_above_mh(energy, mh)
    points = expand_above_ml(energy, ml, points)
    points = expand_above_zm(czr, zm, points, int(zm_expand_limit/hop_len))
    points = trim_points(points)
    # draw_only_wav(sr, signals, t, energy, czr, points, title='Ml = ' + str(ml))
    # draw_a_canvas(sr, signals, t, energy, czr, points, ml, zm)
    draw(sr, signals, t, energy, czr, points)


# 改进后的端点检测方法（用于连读）
def vad2():
    mh = 0.7
    ml = 0.45
    zm = 0.15
    ze = 0.4
    zm_expand_limit = 0.025
    min_speech_time = 0.1
    wav_file = '/Users/d/Project/python_project/speech/wavs/connected.wav'
    window_len = 0.03
    hop_len = 0.005
    zcr_threshold = 100
    sr, signals = siowav.read(wav_file)
    signals = signals / np.max(signals)
    energy, t = st_energy(wav_file, window_len, hop_len, hamming_window)
    czr, t = st_zcr(wav_file, window_len, hop_len, zcr_threshold, hamming_window)
    points = get_above_mh(energy, mh)
    points = expand_above_ml(energy, ml, points)
    points = expand_above_zm(czr, zm, points, int(zm_expand_limit / hop_len))
    points = trim_points(points)
    z_points = get_above_ze(czr, ze)
    points = split_points(points, z_points, min_part_len=min_speech_time/hop_len)
    # draw_only_wav(sr, signals, t, energy, czr, points, title='Ml = ' + str(ml))
    # draw_a_canvas(sr, signals, t, energy, czr, points, ml, zm)
    draw(sr, signals, t, energy, czr, points)


if __name__ == '__main__':
    # vad2()
    vad()


